module Trie (Trie(..),
             empty,
             isEmpty,
             insert,
             lkup,
             remove,
             mapT,
             mapTK,
             filterT,
             filterTK,
             foldT,
             foldTK,
             insertTrie,
             fromList,
             asList,
             dom,
             ran,
             singleton,
             update,
             restrict,
             removeSet,
             pointwise,
             pointwiseRan,
             pointwiseDom,
             mutually,
             mutuallyRan,
             mutuallyDom) where

import AList
import POrd
import Cpo

import Data.Maybe
import qualified Data.Set as S


-- Basic Trie functionality (problem set 1).

data Trie a = Node (Maybe a) (AList (Trie a)) deriving (Eq, Ord)


isEmpty :: Trie a -> Bool
isEmpty (Node Nothing ts) = isNil ts
isEmpty _                 = False


empty :: Trie a
empty = Node Nothing nil


insert :: String -> a -> Trie a -> Trie a
insert "" v (Node x ts)     = Node (Just v) ts
insert (c:cs) v (Node x ts) = case get c ts of
                                Nothing -> let t' = insert cs v empty
                                           in Node x (put c t' ts)
                                Just t  -> let t' = insert cs v t
                                           in Node x (put c t' ts)


lkup :: String -> Trie a -> Maybe a
lkup "" (Node x ts)     = x
lkup (c:cs) (Node x ts) = case get c ts of
                            Nothing -> Nothing
                            Just t  -> lkup cs t


remove :: String -> Trie a -> Trie a
remove "" (Node x ts)     = Node Nothing ts
remove (c:cs) (Node x ts) = case get c ts of
                              Nothing -> Node x ts
                              Just t  -> let t' = remove cs t
                                         in if isEmpty t'
                                            then Node x (del c ts)
                                            else Node x (put c t' ts)


instance Show a => Show (Trie a) where
    show t | isEmpty t = "{}"
    show t | otherwise = "{" ++ (tail $ foldTK (\ k v s -> s ++ "," ++ k ++ "->" ++ show v) "" t) ++ "}"


-- Higher order functions for Tries (problem set 2)

mapT :: (a -> b) -> Trie a -> Trie b
mapT f = mapTK (\ _ v -> f v)


mapTK :: (String -> a -> b) -> Trie a -> Trie b
mapTK f (Node v ts) = mapKV' "" f (Node v ts)
    where
      mapKV' pre f (Node v ts) = let v'  = case v of
                                             Nothing -> Nothing
                                             Just w  -> Just (f pre w)
                                     ts' = map (\ (k, t) -> (k, mapKV' (pre ++ [k]) f t)) ts
                                 in Node v' ts'


filterT :: (a -> Bool) -> Trie a -> Trie a
filterT p = filterTK (\ _ v -> p v)


filterTK :: (String -> a -> Bool) -> Trie a -> Trie a
filterTK p t = filterTK' "" p t
    where
      filterTK' pre p (Node v ts) = let v'  = case v of
                                                 Nothing -> Nothing
                                                 Just w  -> if p pre w then v else Nothing
                                        ts'  = map (\ (k, t) -> (k, filterTK' (pre ++ [k]) p t)) ts
                                        ts'' = filter (not . isEmpty . snd) ts'
                                    in Node v' ts''


foldT :: (a -> b -> b) -> b -> Trie a -> b
foldT f = foldTK (\ _ v a -> f v a)


foldTK :: (String -> a -> b -> b) -> b -> Trie a -> b
foldTK f a = foldTK' "" f a
    where
      foldTK' pre f a (Node Nothing ts)  = foldr (\ (k, t) r -> foldTK' (pre ++ [k]) f r t) a ts
      foldTK' pre f a (Node (Just v) ts) = let tfold = foldr (\ (k, t) r -> foldTK' (pre ++ [k]) f r t) a ts
                                           in f pre v tfold


insertTrie :: Trie a -> Trie a -> Trie a
insertTrie t1 t2 = foldTK insert t1 t2


asList :: Trie a -> [(String, a)]
asList t = foldTK (\ k v l -> (k, v):l) [] t


fromList :: [(String, a)] -> Trie a
fromList = foldr (\ (k, v) t -> insert k v t) empty


zipT :: (Maybe a -> Maybe a -> a) -> Trie a -> Trie a -> Trie a
zipT op t1 t2 = let keys = dom t1 `S.union` dom t2
               in S.fold insertF empty keys
    where
      insertF k t = let v' = (lkup k t1) `op` (lkup k t2)
                    in insert k v' t


-- Map functions for tries (problem set 4)

member :: String -> Trie a  -> Bool
member k t = case lkup k t of { Nothing -> False; Just _ -> True }

singleton :: String -> a -> Trie a
singleton k a = insert k a empty

dom :: Trie a -> S.Set String
dom = foldTK (\ k _ s -> S.insert k s) S.empty

ran :: Ord a => Trie a -> S.Set a
ran = foldT S.insert S.empty

update :: Trie a -> Trie a -> Trie a
update = insertTrie

removeSet :: S.Set String -> Trie a -> Trie a
removeSet s t = S.fold remove t s

restrict  :: Trie a -> S.Set String -> Trie a
restrict t s = filterTK (\ k _ -> S.member k s) t


-- Predicates on Tries (problem set 4).

pointwise :: ((String, a) -> Bool) -> Trie a -> Bool
pointwise p = foldT (&&) True . mapTK (curry p)


pointwiseRan :: (a -> Bool) -> Trie a -> Bool
pointwiseRan p = pointwise (p . snd)


pointwiseDom :: (String -> Bool) -> Trie a -> Bool
pointwiseDom p = pointwise (p . fst)


mutually :: ((String, a) -> (String, a) -> Bool) -> Trie a -> Bool
mutually p = mutually' . asList
    where
      mutually' []        = True
      mutually' (kv1:kvs) = foldr (\ kv2 r -> r && p kv1 kv2) True kvs && mutually' kvs


mutuallyRan :: (a -> a -> Bool) -> Trie a -> Bool
mutuallyRan p = mutually (\ (_, v1) (_, v2) -> p v1 v2)


mutuallyDom :: (String -> String -> Bool) -> Trie a -> Bool
mutuallyDom p = mutually (\ (k1, _) (k2, _) -> p k1 k2)


-- Instance definition for POrd, Cpo (problem set 8).
instance POrd a => POrd (Trie a) where
    m1 `ple` m2 = dom m1 `S.isSubsetOf` dom m2
                  && pointwise (\ (k, v) -> v `ple` fromJust (lkup k m2)) m1

instance Cpo a => Cpo (Trie a) where
    bottom      = empty
    m1 `lub` m2 = let keys = dom m1 `S.union` dom m2
                  in S.fold insertF empty keys
        where
          insertF k m = let v' = fromJust (lkup k m1) `lub` fromJust (lkup k m2)
                        in insert k v' m
